import pygame
import random
import time
from enum import Enum
import sys

# 初始化pygame
pygame.init()
pygame.mixer.init()

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
YELLOW = (255, 255, 0)
GRAY = (128, 128, 128)

# 游戏常量
GRID_SIZE = 20
GRID_WIDTH = 20
GRID_HEIGHT = 20
SCREEN_WIDTH = GRID_WIDTH * GRID_SIZE
SCREEN_HEIGHT = GRID_HEIGHT * GRID_SIZE + 60  # 额外空间用于显示分数
FPS = 10

# 游戏状态枚举
class GameState(Enum):
    MENU = 0
    PLAYING = 1
    GAME_OVER = 2
    SETTINGS = 3

# 难度级别
class Difficulty(Enum):
    EASY = 0.1
    MEDIUM = 0.05
    HARD = 0.02

# 方向枚举
class Direction(Enum):
    UP = (0, -1)
    DOWN = (0, 1)
    LEFT = (-1, 0)
    RIGHT = (1, 0)

# 按钮类
class Button:
    def __init__(self, x, y, width, height, text, color, hover_color):
        self.rect = pygame.Rect(x, y, width, height)
        self.text = text
        self.color = color
        self.hover_color = hover_color
        self.is_hovered = False
        self.font = pygame.font.SysFont('Arial', 20)
        
    def draw(self, surface):
        color = self.hover_color if self.is_hovered else self.color
        pygame.draw.rect(surface, color, self.rect)
        pygame.draw.rect(surface, BLACK, self.rect, 2)
        
        text_surface = self.font.render(self.text, True, BLACK)
        text_rect = text_surface.get_rect(center=self.rect.center)
        surface.blit(text_surface, text_rect)
        
    def check_hover(self, pos):
        self.is_hovered = self.rect.collidepoint(pos)
        return self.is_hovered
        
    def is_clicked(self, pos, event):
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            return self.rect.collidepoint(pos)
        return False

# 食物类
class Food:
    def __init__(self):
        self.position = (0, 0)
        self.color = RED
        self.randomize_position()
        
    def randomize_position(self, snake=None):
        if snake:
            # 确保食物不会生成在蛇身上
            while True:
                self.position = (
                    random.randint(0, GRID_WIDTH - 1),
                    random.randint(0, GRID_HEIGHT - 1)
                )
                if self.position not in snake.positions:
                    break
        else:
            self.position = (
                random.randint(0, GRID_WIDTH - 1),
                random.randint(0, GRID_HEIGHT - 1)
            )
        
    def draw(self, surface):
        rect = pygame.Rect(
            self.position[0] * GRID_SIZE,
            self.position[1] * GRID_SIZE,
            GRID_SIZE,
            GRID_SIZE
        )
        pygame.draw.rect(surface, self.color, rect)
        pygame.draw.rect(surface, BLACK, rect, 1)

# 蛇类
class Snake:
    def __init__(self):
        self.positions = [(GRID_WIDTH // 2, GRID_HEIGHT // 2)]
        self.direction = Direction.RIGHT
        self.next_direction = Direction.RIGHT
        self.color = GREEN
        self.head_color = BLUE
        self.length = 1
        self.score = 0
        self.grow_to = 1
        
    def get_head_position(self):
        return self.positions[0]
        
    def update(self):
        # 更新方向
        self.direction = self.next_direction
        
        # 获取当前头部位置
        head_x, head_y = self.get_head_position()
        
        # 计算新的头部位置
        dir_x, dir_y = self.direction.value
        new_x = (head_x + dir_x) % GRID_WIDTH
        new_y = (head_y + dir_y) % GRID_HEIGHT
        new_head = (new_x, new_y)
        
        # 插入新的头部
        self.positions.insert(0, new_head)
        
        # 如果不需要增长，移除尾部
        if len(self.positions) > self.grow_to:
            self.positions.pop()
        else:
            self.grow_to += 1
        
    def grow(self):
        self.grow_to += 1
        self.score += 10
        
    def change_direction(self, new_direction):
        # 防止180度转弯
        if (new_direction == Direction.UP and self.direction != Direction.DOWN) or \
           (new_direction == Direction.DOWN and self.direction != Direction.UP) or \
           (new_direction == Direction.LEFT and self.direction != Direction.RIGHT) or \
           (new_direction == Direction.RIGHT and self.direction != Direction.LEFT):
            self.next_direction = new_direction
        
    def draw(self, surface):
        for i, pos in enumerate(self.positions):
            rect = pygame.Rect(
                pos[0] * GRID_SIZE,
                pos[1] * GRID_SIZE,
                GRID_SIZE,
                GRID_SIZE
            )
            color = self.head_color if i == 0 else self.color
            pygame.draw.rect(surface, color, rect)
            pygame.draw.rect(surface, BLACK, rect, 1)
            
    def check_collision(self):
        # 检查是否撞到自己
        return self.get_head_position() in self.positions[1:]

# 游戏类
class Game:
    def __init__(self):
        self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
        pygame.display.set_caption("Snake Game")
        self.clock = pygame.time.Clock()
        self.font = pygame.font.SysFont('Arial', 24)
        self.big_font = pygame.font.SysFont('Arial', 36)
        
        self.state = GameState.MENU
        self.difficulty = Difficulty.MEDIUM
        self.high_score = 0
        self.paused = False
        
        # 创建游戏对象
        self.snake = Snake()
        self.food = Food()
        
        # 创建按钮
        self.start_button = Button(SCREEN_WIDTH//2 - 100, SCREEN_HEIGHT//2 - 30, 200, 50, "Start Game", GREEN, YELLOW)
        self.settings_button = Button(SCREEN_WIDTH//2 - 100, SCREEN_HEIGHT//2 + 30, 200, 50, "Settings", BLUE, YELLOW)
        self.exit_button = Button(SCREEN_WIDTH//2 - 100, SCREEN_HEIGHT//2 + 90, 200, 50, "Exit", RED, YELLOW)
        
        self.resume_button = Button(SCREEN_WIDTH//2 - 100, SCREEN_HEIGHT//2 - 30, 200, 50, "Resume", GREEN, YELLOW)
        self.menu_button = Button(SCREEN_WIDTH//2 - 100, SCREEN_HEIGHT//2 + 30, 200, 50, "Main Menu", BLUE, YELLOW)
        
        self.restart_button = Button(SCREEN_WIDTH//2 - 100, SCREEN_HEIGHT//2, 200, 50, "Restart", GREEN, YELLOW)
        self.game_over_menu_button = Button(SCREEN_WIDTH//2 - 100, SCREEN_HEIGHT//2 + 60, 200, 50, "Main Menu", BLUE, YELLOW)
        
        self.easy_button = Button(SCREEN_WIDTH//2 - 100, SCREEN_HEIGHT//2 - 60, 200, 50, "Easy", GREEN, YELLOW)
        self.medium_button = Button(SCREEN_WIDTH//2 - 100, SCREEN_HEIGHT//2, 200, 50, "Medium", BLUE, YELLOW)
        self.hard_button = Button(SCREEN_WIDTH//2 - 100, SCREEN_HEIGHT//2 + 60, 200, 50, "Hard", RED, YELLOW)
        self.back_button = Button(SCREEN_WIDTH//2 - 100, SCREEN_HEIGHT//2 + 120, 200, 50, "Back", GRAY, YELLOW)
        
    def reset_game(self):
        self.snake = Snake()
        self.food.randomize_position(self.snake)
        self.paused = False
        
    def handle_events(self):
        mouse_pos = pygame.mouse.get_pos()
        
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False
                
            if self.state == GameState.MENU:
                self.start_button.check_hover(mouse_pos)
                self.settings_button.check_hover(mouse_pos)
                self.exit_button.check_hover(mouse_pos)
                
                if self.start_button.is_clicked(mouse_pos, event):
                    self.state = GameState.PLAYING
                    self.reset_game()
                elif self.settings_button.is_clicked(mouse_pos, event):
                    self.state = GameState.SETTINGS
                elif self.exit_button.is_clicked(mouse_pos, event):
                    return False
                    
            elif self.state == GameState.PLAYING:
                if event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_UP:
                        self.snake.change_direction(Direction.UP)
                    elif event.key == pygame.K_DOWN:
                        self.snake.change_direction(Direction.DOWN)
                    elif event.key == pygame.K_LEFT:
                        self.snake.change_direction(Direction.LEFT)
                    elif event.key == pygame.K_RIGHT:
                        self.snake.change_direction(Direction.RIGHT)
                    elif event.key == pygame.K_p:
                        self.paused = not self.paused
                    elif event.key == pygame.K_ESCAPE:
                        self.state = GameState.MENU
                        
            elif self.state == GameState.GAME_OVER:
                self.restart_button.check_hover(mouse_pos)
                self.game_over_menu_button.check_hover(mouse_pos)
                
                if self.restart_button.is_clicked(mouse_pos, event):
                    self.state = GameState.PLAYING
                    self.reset_game()
                elif self.game_over_menu_button.is_clicked(mouse_pos, event):
                    self.state = GameState.MENU
                    
            elif self.state == GameState.SETTINGS:
                self.easy_button.check_hover(mouse_pos)
                self.medium_button.check_hover(mouse_pos)
                self.hard_button.check_hover(mouse_pos)
                self.back_button.check_hover(mouse_pos)
                
                if self.easy_button.is_clicked(mouse_pos, event):
                    self.difficulty = Difficulty.EASY
                elif self.medium_button.is_clicked(mouse_pos, event):
                    self.difficulty = Difficulty.MEDIUM
                elif self.hard_button.is_clicked(mouse_pos, event):
                    self.difficulty = Difficulty.HARD
                elif self.back_button.is_clicked(mouse_pos, event):
                    self.state = GameState.MENU
                    
        return True
        
    def update(self):
        if self.state == GameState.PLAYING and not self.paused:
            self.snake.update()
            
            # 检查是否吃到食物
            if self.snake.get_head_position() == self.food.position:
                self.snake.grow()
                self.food.randomize_position(self.snake)
                
                # 更新最高分
                if self.snake.score > self.high_score:
                    self.high_score = self.snake.score
                    
            # 检查碰撞
            if self.snake.check_collision():
                self.state = GameState.GAME_OVER
                
    def draw(self):
        self.screen.fill(WHITE)
        
        if self.state == GameState.MENU:
            # 绘制菜单界面
            title = self.big_font.render("Snake Game", True, BLACK)
            title_rect = title.get_rect(center=(SCREEN_WIDTH//2, SCREEN_HEIGHT//4))
            self.screen.blit(title, title_rect)
            
            self.start_button.draw(self.screen)
            self.settings_button.draw(self.screen)
            self.exit_button.draw(self.screen)
            
        elif self.state == GameState.PLAYING:
            # 绘制游戏区域
            for x in range(GRID_WIDTH):
                for y in range(GRID_HEIGHT):
                    rect = pygame.Rect(x * GRID_SIZE, y * GRID_SIZE, GRID_SIZE, GRID_SIZE)
                    pygame.draw.rect(self.screen, (230, 230, 230), rect)
                    pygame.draw.rect(self.screen, (200, 200, 200), rect, 1)
                    
            # 绘制蛇和食物
            self.snake.draw(self.screen)
            self.food.draw(self.screen)
            
            # 绘制分数
            score_text = self.font.render(f"Score: {self.snake.score}", True, BLACK)
            high_score_text = self.font.render(f"Highest Score: {self.high_score}", True, BLACK)
            self.screen.blit(score_text, (10, SCREEN_HEIGHT - 50))
            self.screen.blit(high_score_text, (10, SCREEN_HEIGHT - 25))
            
            # 绘制暂停按钮
            if self.paused:
                pause_text = self.big_font.render("Pause", True, BLACK)
                pause_rect = pause_text.get_rect(center=(SCREEN_WIDTH//2, SCREEN_HEIGHT//2))
                self.screen.blit(pause_text, pause_rect)
                
                self.resume_button.draw(self.screen)
                self.menu_button.draw(self.screen)
                
        elif self.state == GameState.GAME_OVER:
            # 绘制Game Over界面
            game_over_text = self.big_font.render("Game Over", True, RED)
            game_over_rect = game_over_text.get_rect(center=(SCREEN_WIDTH//2, SCREEN_HEIGHT//4))
            self.screen.blit(game_over_text, game_over_rect)
            
            score_text = self.font.render(f"Your Score: {self.snake.score}", True, BLACK)
            high_score_text = self.font.render(f"Highest Score: {self.high_score}", True, BLACK)
            score_rect = score_text.get_rect(center=(SCREEN_WIDTH//2, SCREEN_HEIGHT//4 + 60))
            high_score_rect = high_score_text.get_rect(center=(SCREEN_WIDTH//2, SCREEN_HEIGHT//4 + 90))
            
            self.screen.blit(score_text, score_rect)
            self.screen.blit(high_score_text, high_score_rect)
            
            self.restart_button.draw(self.screen)
            self.game_over_menu_button.draw(self.screen)
            
        elif self.state == GameState.SETTINGS:
            # 绘制Settings界面
            settings_text = self.big_font.render("Settings", True, BLACK)
            settings_rect = settings_text.get_rect(center=(SCREEN_WIDTH//2, SCREEN_HEIGHT//4))
            self.screen.blit(settings_text, settings_rect)
            
            difficulty_text = self.font.render("Select Difficulty:", True, BLACK)
            difficulty_rect = difficulty_text.get_rect(center=(SCREEN_WIDTH//2, SCREEN_HEIGHT//4 + 60))
            self.screen.blit(difficulty_text, difficulty_rect)
            
            self.easy_button.draw(self.screen)
            self.medium_button.draw(self.screen)
            self.hard_button.draw(self.screen)
            self.back_button.draw(self.screen)
            
        pygame.display.flip()
        
    def run(self):
        running = True
        while running:
            running = self.handle_events()
            self.update()
            self.draw()
            
            # 根据难度调整游戏速度
            if self.state == GameState.PLAYING and not self.paused:
                self.clock.tick(FPS + int(self.snake.score * self.difficulty.value))
            else:
                self.clock.tick(FPS)
                
        pygame.quit()
        sys.exit()

# 启动游戏
if __name__ == "__main__":
    game = Game()
    game.run()
